

sensPlot = function(x, 
                     contour.levels = NULL,
                     col.zero = "blue",
                     lty.zero = 1,
                     col.insig = "blue",
                     lty.insig = 1,
                     data.line = TRUE,
                     X.pch = NULL,  #vector of length 3: non-transformed points, transformed points, points outside plotting range
                     signif.level = 0.05,
                     labcex = 0.75,
                     limit.Xplot = FALSE, #MH: limit plotting covariates to enlarge contour
                     txtlab = FALSE,  #add text label to the plots of covariates.
                     which.txtlab = NULL, #enter numeric vector to specify which label to show. e.g. c(1:3) shows first 3 covariates.
                     ...) {
    #note in help: if contours are too rough, up nsim in sens fn
    if(class(x) == "sensitivity"){  
        sensPlotMain(x, contour.levels, col.zero, lty.zero, col.insig, lty.insig, data.line, X.pch, signif.level, labcex, limit.Xplot, txtlab, which.txtlab,...)
    }else if(class(x) == "sensitivityCombo"){
        sensPlotCombo(x, contour.levels, col.zero, lty.zero, col.insig, lty.insig, data.line, X.pch, signif.level, labcex, limit.Xplot, txtlab, which.txtlab,...)
    }else{
        stop("x must be of class sensitivity or sensitivity.combo")
    }  
} #end of sensPlot

assignInNamespace("sensPlot",sensPlot, ns="treatSens") 




sensPlotMain = function(x, contour.levels, col.zero, lty.zero, col.insig, lty.insig, data.line, X.pch, signif.level, labcex, limit.Xplot, txtlab, which.txtlab,...){
    ##Add row/column for zeta = 0 if not included in grid
    null.tau=x$tau0
    null.se=x$se.tau0
    part.cors = x$sensParam == "cor"
    
    Zcors = as.numeric(dimnames(x$sp.z)[[2]]) #horizontal grids of U
    Ycors = as.numeric(dimnames(x$sp.y)[[1]]) #vertical grids of U
    
    if(part.cors){
        print.text <- "Partial correlations with U"
    }else {
        print.text <- "Coefficients on U" 
    }
    
    ##############  
    #######Only include X on the plot for GLM-style
    ############
    Xpart = x$Xcoef[!is.na(x$Xcoef[,1]) & !is.na(x$Xcoef[,2]),] #coefficients of null model.  
    Xpart.plot = x$Xcoef.plot[!is.na(x$Xcoef.plot[,1]) & !is.na(x$Xcoef.plot[,2]),]
    if(!is.null(Xpart) & is.null(dim(Xpart))){
        Xpart = matrix(Xpart, nrow = 1)
        Xpart.plot = matrix(Xpart.plot, nrow = 1)
        
    } 
    
    Xpart.plot2 = cbind(Xpart.plot[,1],Xpart.plot[,2], ifelse(Xpart[,2]>=0,1,2)) #MH: add sign of coef of X on Y to Xpart  
    if(!is.null(Xpart.plot2) & is.null(dim(Xpart.plot2))){
        Xpart.plot2 = matrix(Xpart.plot2, nrow = 1)
    } 
    
    #note that due to correlation among Xs, some may not appear on plot
    #because observed partial cors don't map directly to coefs in this case
    #forcing inclusion can lead to difficult to read plot.  
    if(is.null(X.pch)){
        out.pch = 1
        X.pch = ifelse(Xpart[,2]>=0,3,6) # plus sign for non-transformed plots, reverse triangle for transformed plots
    }else{
        out.pch = X.pch[3]
        X.pch = ifelse(Xpart[,2]>=0,X.pch[1],X.pch[2])	
    }
    if (any(Xpart[,2]<0)) {
        cat("Note: Predictors with negative coefficients for the response surface have been transformed through multiplication by -1 and are displayed as inverted triangles.", "\n")    
    }
    
    
    nr = length(Zcors); nc = length(Ycors)
    taus.est = t(apply(x$tau, c(1,2), mean, na.rm = T))
    K = dim(x$se.tau)[3]
    W = apply(x$se.tau^2, c(1,2), mean, na.rm = T)
    B = apply(x$tau, c(1,2), sd, na.rm = T)^2
    se.est = t(sqrt(W+(1+1/K)*B))
    taus = taus.est
    se.taus = se.est
    taus[Zcors==0,] = null.tau
    taus[,Ycors==0] = null.tau
    
    dimnames(taus)[[1]] = Zcors
    dimnames(taus)[[2]] = Ycors
    taus <<- taus ## is this right? it potentially changes the value of tau in a parent environment
    dimnames(se.taus)[[1]] = Zcors
    dimnames(se.taus)[[2]] = Ycors
    se.taus <<- se.taus
    
    if(is.null(contour.levels)){
        exTau = c(taus[dim(taus)[1], dim(taus)[2]], taus[1,dim(taus)[2]]) #extreme values of tau at right end
        clevels = round(seq(exTau[2]*.8, exTau[1]*.8, length.out = 14), 2) #vals at which contours are drawn
    }else{
        clevels = contour.levels
    } 
    
    par(mgp = c(2,.5,0)) #dist of axis label, tick mark label, tick mark
    if(part.cors){
        xlab = expression(paste("Partial cor. with U in model for treatment, ", rho^zu))
        ylab = expression(paste("Partial cor. with U in model for response, ", rho^yu))
    }else{
        #xlab = expression(paste("Coef. on U in model for treatment, ", zeta^z))
        #ylab = expression(paste("Coef. on U in model for response, ", zeta^y))
        xlab = "Effect on Treatment" #vb-edit, changes the default axis text
        ylab = "Effect on Outcome" #vb-edit, changes the default axis text
    }
    
    if (limit.Xplot) {  
        #old codes
        plot(Xpart.plot2[,1], Xpart.plot2[,2], col=c("red","blue")[Xpart.plot2[,3]], xlim = c(min(Zcors, na.rm = T),max(Zcors, na.rm = T)), 
             ylim = c(min(Ycors, na.rm = T),max(Ycors, na.rm = T)), pch = X.pch, xlab = xlab, ylab = ylab)

        outsidePts = Xpart.plot2[(Xpart.plot[,1] < min(Zcors, na.rm = T)) | (Xpart.plot[,1] > max(Zcors, na.rm = T)) | (Xpart.plot[,2] > max(Ycors, na.rm = T)),]
        if(length(outsidePts) > 0){
            if(is.null(dim(outsidePts))){
                outsidePts[1] = max(outsidePts[1], min(Zcors, na.rm = T))
                outsidePts[1] = min(outsidePts[1], max(Zcors, na.rm = T))
                outsidePts[2] = min(outsidePts[2], max(Ycors, na.rm = T))
                outsidePts = matrix(outsidePts, nrow = 1)
            }else{
                outsidePts[,1] = apply(cbind(outsidePts[,1], min(Zcors, na.rm = T)), 1, max)
                outsidePts[,1] = apply(cbind(outsidePts[,1], max(Zcors, na.rm = T)), 1, min)
                outsidePts[,2] = apply(cbind(outsidePts[,2], max(Ycors, na.rm = T)), 1, min)
            }
            points(outsidePts[,1], outsidePts[,2], col=c("red","blue")[outsidePts[,3]], pch = out.pch, cex = 1.5, lwd = 3)

            warning("Note: predictors outside plot region plotted at margin with O")
        }
    } else {
        #MH: define max, min of plots
        xplot.min = ifelse(min(Zcors, na.rm = T)<min(Xpart.plot[,1]),min(Zcors, na.rm = T),min(Xpart.plot[,1]))
        xplot.max = ifelse(max(Zcors, na.rm = T)>max(Xpart.plot[,1]),max(Zcors, na.rm = T),max(Xpart.plot[,1]))
        yplot.max = ifelse(max(Ycors, na.rm = T)>max(Xpart.plot[,2]),max(Ycors, na.rm = T),max(Xpart.plot[,2]))  
        #plot(Xpart.plot2[,1], Xpart.plot2[,2], col=c("red","blue")[as.factor(Xpart.plot2[,3])], xlim = c(xplot.min,xplot.max),
        #     ylim = c(0,yplot.max), pch = X.pch, xlab = xlab, ylab = ylab)
        plot(Xpart.plot2[,1], Xpart.plot2[,2], col=c("black","black")[as.factor(Xpart.plot2[,3])], xlim = c(-3,3),
             ylim = c(0,1), pch = X.pch, xlab = xlab, ylab = ylab)  #vb-edit, makes the crosses and triangles black, max out y at 1
        
    }
    
    #codes for txtlab
    if (txtlab) {
        if (is.null(which.txtlab)) { #show all text label
            text(Xpart.plot2[,1], Xpart.plot2[,2], labels=x$varnames[-c(1,2)], cex=labcex, pos=1)
        } else { #show selected label
            which.txtlab2 = which.txtlab + 2
            varnames2 = x$varnames
            varnames2[as.numeric(paste(- which.txtlab2))] = ""
            text(Xpart.plot2[,1], Xpart.plot2[,2], labels=varnames2[-c(1,2)], cex=labcex, pos=1)
        }
    }
    
    #abline(h = 0)  #vb-edit, remove
    #abline(v = 0)  #vb-edit, remove

    
    legend(0.8*max(Zcors), 0, legend = round(x$tau0,2), cex = labcex,
           yjust = 0.5, x.intersp = 0, y.intersp = 0,
           bg = ifelse(par("bg")== "transparent", "white", par("bg")), 
           box.lty = 0)
    
    box()
    
    #    contour(Zcors, Ycors, taus, levels = clevels, 
    #add = T, labcex = labcex,...)
    contour(Zcors, Ycors, taus, levels = clevels, 
            add = T, labcex = labcex, col="lightgray",...)  #vb-edit, makes the contour lines lighter color
    
    
    
    contour(Zcors, Ycors, taus, levels = 0, lwd = 2,
            add = T, col = col.zero,lty = lty.zero,labcex = labcex,...)
    
    contour(as.numeric(dimnames(x$sp.z)[[2]]), as.numeric(dimnames(x$sp.y)[[1]]), taus.est/se.est, labels = "N.S.",
            #levels = c(-1, 1)*qnorm(signif.level/2), add = T, col = col.insig,
            levels = c(-1, 1)*qnorm(signif.level/2), add = T, col = "black",  #vb-edit
            lty = lty.insig, labcex = labcex, lwd = 2,...)
    
    if (all(sign(Xpart[,1])!=sign(x$tau0))){
        warning("Cannot add data line because XXXXX.")
    }else{
        if(data.line & length(Xpart)>1){
            proj.pts = apply(Xpart.plot^2, 1, mean)
            max.pt = Xpart.plot[proj.pts == max(proj.pts),]
            max.pt[1] = sign(x$tau0)*abs(max.pt[1])
            Zdiff = (Zcors-max.pt[1])*sign(max.pt[1])
            Zgrt = which(Zdiff < 0)
            zcor = Zgrt[which(Zdiff[Zgrt] == max(Zdiff[Zgrt]))]  
            if((Zcors[zcor] > max.pt[1] & zcor > 1)||(zcor==length(Zcors))){ 
                zpts = c(zcor-1, zcor)
            }else{
                zpts = c(zcor, zcor+1)
            }
            Ydiff = Ycors-max.pt[2]
            Ygrt = which(Ydiff < 0)
            ycor = Ygrt[which(Ydiff[Ygrt] == max(Ydiff[Ygrt]))]
            if((Ycors[ycor] > max.pt[2] & ycor > 1)||(ycor==length(Ycors))){ 
                ypts = c(ycor-1, ycor)
            }else{
                ypts = c(ycor, ycor+1)
            }
            clevel = ((Zcors[zpts[2]] - Zcors[zpts[1]])*(Ycors[ypts[2]] - Ycors[ypts[1]]))^(-1)*
                sum(taus[zpts, ypts]*
                        matrix(c(-(Zcors[zpts[2]] - max.pt[1])*(Ycors[ypts[1]] - max.pt[2]), 
                                 (Zcors[zpts[1]] - max.pt[1])*(Ycors[ypts[1]] - max.pt[2]),
                                 (Zcors[zpts[2]] - max.pt[1])*(Ycors[ypts[2]] - max.pt[2]),
                                 -(Zcors[zpts[1]] - max.pt[1])*(Ycors[ypts[2]] - max.pt[2])), 
                               nrow = 2, byrow = T))
            contour(Zcors, Ycors, taus, levels = round(clevel,2),
                    add = T, col = "grey",labcex = labcex, lwd = 2,...)
        }else{
            if(data.line)
                warning("Cannot add data line because there are no non-treatment covariates.")
        }
    }

}




assignInNamespace("sensPlot",sensPlot, ns="treatSens") 
